#include "mt3d_cuda.h"

void mt3d::post_process()
{
    std::clog << "Processing MT stations ... \n";

    // Sort out MT stations
    mt_station tmp_station;

    for (size_t i = 0; i < node_num_; i++)
    {
        if (node_markers_[i] == station_marker_)
        {
            tmp_station.vert = mesh_nodes_.get(i);
            obs_site_.push_back(tmp_station);
        }
    }

    obs_num_ = obs_site_.size();

    // Find tetrahedrons that include mt stations
    for (size_t i = 0; i < elem_num_; i++)
    {
        if (elem_markers_[i] != air_marker_) // exclude the air layer to obtain smoother results
        {
            for (size_t j = 0; j < 4; j++)
            {
                if (node_markers_[mesh_elems_[i].vert[j]->id] == station_marker_)
                {
                    for (size_t k = 0; k < obs_num_; k++)
                    {
                        if (obs_site_[k].vert == mesh_elems_[i].vert[j])
                        {
                            obs_site_[k].host.push_back(mesh_elems_.get(i));
                            break;
                        }
                    }
                }
            }
        }
    }

    // Calculate observations
    gctl::array<double> total_vol(obs_num_, 0.0);
    for (size_t i = 0; i < obs_num_; i++)
    {
        for (size_t j = 0; j < obs_site_[i].host.size(); j++)
        {
            total_vol[i] += obs_site_[i].host[j]->volume();
        }
    }

    efem_sf::ordertype_e o_type;
    gctl::vertex3dc* tmp_vert[4];
    double ux, uy, uz, vx, vy, vz, dummy, tmp_vol;
    
    cuda_cd te_ex, te_hy, tm_ey, tm_hx;
    gctl::array<cuda_cd> TE_EX(obs_num_, _zero), TE_HY(obs_num_, _zero), TM_EY(obs_num_, _zero), TM_HX(obs_num_, _zero);

    for (size_t i = 0; i < obs_num_; i++)
    {
        for (size_t j = 0; j < obs_site_[i].host.size(); j++)
        {
            tmp_vert[0] = obs_site_[i].host[j]->vert[0];
            tmp_vert[1] = obs_site_[i].host[j]->vert[1];
            tmp_vert[2] = obs_site_[i].host[j]->vert[2];
            tmp_vert[3] = obs_site_[i].host[j]->vert[3];

            te_ex = te_hy = tm_ey = tm_hx = _zero;
            for (size_t d = 0; d < 6; d++)
            {
                if (elem_edge_reversed_[obs_site_[i].host[j]->id][d]) o_type = efem_sf::Reverse;
                else o_type = efem_sf::Normal;

                elsf_.tetrahedron(obs_site_[i].vert->x, obs_site_[i].vert->y, obs_site_[i].vert->z, 
                    tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, 
                    d, efem_sf::Value, o_type, ux, uy, uz, dummy, dummy, dummy);

                elsf_.tetrahedron(obs_site_[i].vert->x, obs_site_[i].vert->y, obs_site_[i].vert->z, 
                    tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, 
                    d, efem_sf::Curl, o_type, vx, vy, vz, dummy, dummy, dummy);

                te_ex = clcg_Zsum(te_ex, clcg_Zscale(uy, h_X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]]));
                te_hy = clcg_Zsum(te_hy, clcg_Zscale(vx, h_X_TE_[elem_edge_index_[obs_site_[i].host[j]->id][d]]));

                tm_ey = clcg_Zsum(tm_ey, clcg_Zscale(ux, h_X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]]));
                tm_hx = clcg_Zsum(tm_hx, clcg_Zscale(vy, h_X_TM_[elem_edge_index_[obs_site_[i].host[j]->id][d]]));
            }

            tmp_vol = obs_site_[i].host[j]->volume();

            TE_EX[i] = clcg_Zsum(TE_EX[i], clcg_Zscale(tmp_vol/total_vol[i], te_ex));
            TE_HY[i] = clcg_Zsum(TE_HY[i], clcg_Zscale(tmp_vol/total_vol[i], te_hy));
            TM_EY[i] = clcg_Zsum(TM_EY[i], clcg_Zscale(tmp_vol/total_vol[i], tm_ey));
            TM_HX[i] = clcg_Zsum(TM_HX[i], clcg_Zscale(tmp_vol/total_vol[i], tm_hx));
        }

        TE_HY[i] = cuCdiv(TE_HY[i], clcg_Zscale(omega_*_mu0, _1i));
        TM_HX[i] = cuCdiv(TM_HX[i], clcg_Zscale(omega_*_mu0, _1i));     
    }

    // Add primary field
    le_.set_to_TE_mode();
    le_.initialize();
    for (size_t i = 0; i < obs_num_; i++)
    {
        le_.calculate_EH(-1.0*obs_site_[i].vert->z, 0.0); // Zero obique angle

        TE_EX[i] = clcg_Zsum(TE_EX[i], lcg2cuda_complex(le_.get_Ex()));
        TE_HY[i] = clcg_Zsum(TE_HY[i], lcg2cuda_complex(le_.get_Hy()));
    }

    le_.set_to_TM_mode();
    le_.initialize();
    for (size_t i = 0; i < obs_num_; i++)
    {
        le_.calculate_EH(-1.0*obs_site_[i].vert->z, 0.0); // Zero obique angle

        TM_EY[i] = clcg_Zsum(TM_EY[i], lcg2cuda_complex(le_.get_Ey()));
        TM_HX[i] = clcg_Zsum(TM_HX[i], lcg2cuda_complex(le_.get_Hx()));
    }

    for (size_t i = 0; i < obs_num_; i++)
    {
        obs_site_[i].Zxy = cuCdiv(TE_EX[i], TE_HY[i]);
        obs_site_[i].Zyx = cuCdiv(TM_EY[i], TM_HX[i]);
    }

    return;
}
